"""
icon_generator.py - 游戏图标生成器
使用g4f生成游戏图标，保存到本地
只需要在有g4f环境的设备上运行一次
"""

import os
import sys
import time
import requests
from g4f.client import Client

class GameIconGenerator:
    """
    游戏图标生成器
    使用g4f生成游戏图标并保存到本地
    """
    
    # 游戏图标定义 - 包含图标名称和对应的提示词
    GAME_ICONS = {
        "road": {
            "filename": "road.png",
            "prompt": "一个纯黑色4×4方块,不是画画，就是一个黑色的图标一个黑色的方块正方形 "
                     "纯黑色，不要有任何添加,也不要渐变色就是一个黑色方块",
            "description": "道路图标"
        },
        "house": {
            "filename": "house.png",
            "prompt": "A simple red house icon with triangle roof, pixel art style, flat design, "
                     "white background, cartoon style, game asset, clean lines",
            "description": "住宅图标"
        },
        "commercial": {
            "filename": "commercial.png",
            "prompt": "A blue office building icon, minimalist tall rectangle with small window squares, "
                     "flat design, white background, corporate building, simple geometry,建筑高度和宽度大约是四比三",
            "description": "商业建筑图标"
        },
        "industrial": {
            "filename": "industrial.png",
            "prompt": "A gray factory icon with chimney, industrial style, flat design, "
                     "white background, minimalist factory building, smoke from chimney",
            "description": "工业建筑图标"
        },
        "park": {
            "filename": "park.png",
            "prompt": "A green park icon with tree and bench, flat design, cartoon minimal style, "
                     "white background, nature icon, recreational area",
            "description": "公园图标"
        },
        "empty": {
            "filename": "empty.png",
            "prompt": "一个纯白色的4×4的方块正方形的一个方块不要画什么就是一个纯白色的方块正方形 "
                     "纯白色里面不要有任何添加，也不需要渐变色",
            "description": "空地图标"
        }
    }
    
    def __init__(self, output_dir="pgame/icons"):
        """
        初始化图标生成器
        
        Args:
            output_dir: 图标输出目录
        """
        self.output_dir = output_dir
        self.client = None
        self.setup_directories()
    
    def setup_directories(self):
        """创建必要的目录结构"""
        os.makedirs(self.output_dir, exist_ok=True)
        print(f"图标将保存到: {os.path.abspath(self.output_dir)}")
    
    def initialize_g4f_client(self):
        """初始化g4f客户端"""
        try:
            self.client = Client()
            print("g4f客户端初始化成功")
            return True
        except Exception as e:
            print(f"g4f客户端初始化失败: {e}")
            return False
    
    def generate_icon(self, icon_name, icon_info, retry_count=3):
        """
        生成单个图标
        
        Args:
            icon_name: 图标名称
            icon_info: 图标信息字典
            retry_count: 重试次数
            
        Returns:
            bool: 是否生成成功
        """
        output_path = os.path.join(self.output_dir, icon_info["filename"])
        
        # 检查图标是否已存在
        if os.path.exists(output_path):
            print(f"✓ 图标 '{icon_name}' ({icon_info['description']}) 已存在，跳过生成")
            return True
        
        print(f"正在生成图标: {icon_name} ({icon_info['description']})...")
        print(f"提示词: {icon_info['prompt']}")
        
        for attempt in range(retry_count):
            try:
                # 调用g4f生成图片
                response = self.client.images.generate(
                    model="flux",
                    prompt=icon_info["prompt"],
                    response_format="url"
                )
                
                image_url = response.data[0].url
                print(f"  生成成功，下载URL: {image_url[:50]}...")
                
                # 下载图片
                headers = {
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
                }
                img_response = requests.get(image_url, headers=headers, timeout=45)
                img_response.raise_for_status()
                
                # 保存图片
                with open(output_path, 'wb') as f:
                    f.write(img_response.content)
                
                # 验证图片文件
                if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
                    print(f"✓ 图标 '{icon_name}' 保存成功: {output_path}")
                    time.sleep(2)  # 避免请求过快
                    return True
                else:
                    print(f"✗ 图标 '{icon_name}' 保存失败，文件为空或不存在")
                    
            except requests.exceptions.RequestException as e:
                print(f"✗ 下载图标 '{icon_name}' 失败 (尝试 {attempt+1}/{retry_count}): {e}")
            except Exception as e:
                print(f"✗ 生成图标 '{icon_name}' 失败 (尝试 {attempt+1}/{retry_count}): {e}")
            
            if attempt < retry_count - 1:
                print(f"  等待5秒后重试...")
                time.sleep(5)
        
        print(f"✗ 图标 '{icon_name}' 生成失败，已达到最大重试次数")
        return False
    
    def generate_all_icons(self):
        """生成所有游戏图标"""
        print("=" * 60)
        print("开始生成游戏图标...")
        print("=" * 60)
        
        # 初始化g4f客户端
        if not self.initialize_g4f_client():
            print("无法初始化g4f客户端，请检查g4f安装和网络连接")
            return False
        
        success_count = 0
        total_icons = len(self.GAME_ICONS)
        
        for icon_name, icon_info in self.GAME_ICONS.items():
            if self.generate_icon(icon_name, icon_info):
                success_count += 1
        
        # 生成结果汇总
        print("=" * 60)
        print("图标生成完成!")
        print(f"成功生成: {success_count}/{total_icons} 个图标")
        print(f"图标保存位置: {os.path.abspath(self.output_dir)}")
        
        if success_count == total_icons:
            print("✓ 所有图标生成成功!")
            print("\n重要提示:")
            print("1. 将整个 'pgame' 目录复制到游戏服务器")
            print("2. 在其他设备上运行游戏时不需要g4f")
            print("3. 游戏将直接使用这些本地图标文件")
            return True
        else:
            print("⚠ 部分图标生成失败，游戏可能无法正常显示")
            print("请检查网络连接后重新运行此脚本")
            return False
    
    def verify_icons(self):
        """验证所有图标文件是否存在"""
        print("验证图标文件...")
        missing_icons = []
        
        for icon_name, icon_info in self.GAME_ICONS.items():
            icon_path = os.path.join(self.output_dir, icon_info["filename"])
            if os.path.exists(icon_path):
                print(f"✓ {icon_info['description']}: 存在")
            else:
                print(f"✗ {icon_info['description']}: 缺失")
                missing_icons.append(icon_info["description"])
        
        if missing_icons:
            print(f"\n缺少 {len(missing_icons)} 个图标: {', '.join(missing_icons)}")
            return False
        else:
            print("✓ 所有图标文件验证通过")
            return True


def main():
    """主函数"""
    print("游戏图标生成器")
    print("说明: 此脚本使用g4f生成游戏所需图标")
    print("只需在有g4f环境的设备上运行一次\n")
    
    generator = GameIconGenerator()
    
    # 检查是否已存在图标
    if generator.verify_icons():
        print("\n所有图标已存在，是否重新生成? (y/N): ", end="")
        choice = input().strip().lower()
        if choice != 'y':
            print("已跳过生成，使用现有图标")
            return
    
    # 生成图标
    print("\n开始生成图标，这可能需要几分钟时间...")
    success = generator.generate_all_icons()
    
    if success:
        print("\n图标生成完成! 您现在可以:")
        print("1. 将 'pgame' 目录复制到其他设备")
        print("2. 运行 city_game_final.py 启动游戏")
    else:
        print("\n图标生成失败，请检查:")
        print("1. g4f是否正确安装: pip install g4f")
        print("2. 网络连接是否正常")
        print("3. g4f服务是否可用")


if __name__ == "__main__":
    main()
